package org.example.kafka;

import org.apache.jmeter.config.Arguments;
import org.apache.jmeter.protocol.java.sampler.AbstractJavaSamplerClient;
import org.apache.jmeter.protocol.java.sampler.JavaSamplerContext;
import org.apache.jmeter.samplers.SampleResult;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerRecord;

import java.util.Properties;

public class KafkaProducerTest extends AbstractJavaSamplerClient {
    private static String MSG_SER = "org.apache.kafka.common.serialization.StringSerializer";
    private static String KEY_SER = "org.apache.kafka.common.serialization.StringSerializer";
    private ProducerRecord<String,String> record;

    @Override
    public void setupTest(JavaSamplerContext context) {

    }

    @Override
    public Arguments getDefaultParameters() {
        Arguments arguments = new Arguments();
        arguments.addArgument("bootstrap.servers", "");
        arguments.addArgument("kafka_topic", "");
        arguments.addArgument("kafka_message", "");
        return arguments;
    }

    public SampleResult runTest(JavaSamplerContext javaSamplerContext) {
        String kafka_brokers = javaSamplerContext.getParameter("bootstrap.servers");
        String kafka_topic = javaSamplerContext.getParameter("kafka_topic");
        String kafka_message = javaSamplerContext.getParameter("kafka_message");
        Properties props = new Properties();
        props.put("bootstrap.servers", kafka_brokers);
        props.put("key.serializer", KEY_SER);
        props.put("value.serializer", MSG_SER);
        props.put("acks", "1");
        KafkaProducer<String, String> producer = new KafkaProducer<>(props);
        record = new ProducerRecord<String, String>(kafka_topic, kafka_message);
        SampleResult result = new SampleResult();
        result.sampleStart();
        try {
            producer.send(record);
            result.setResponseHeaders("response: " + "消息发送成功");
            result.setRequestHeaders("request: " + kafka_message);
            result.setDataType(SampleResult.TEXT);
            result.setSuccessful(true);
        }catch (Exception e){
            e.printStackTrace();
        }finally {
            result.sampleEnd();
            producer.close();
        }

        return null;
    }